import time as timer
import heapq
import random
from single_agent_planner import compute_heuristics, a_star, get_location, get_sum_of_cost


def detect_collision(path1, path2):
    # 检测两条路径之间的第一个碰撞（顶点碰撞或边碰撞）
    # 如果没有碰撞，返回 None
    pair_loc1 = [None] * 2
    pair_loc2 = [None] * 2
    for ts in range(max(len(path1), len(path2))):
        pair_loc1[0] = get_location(path1, ts)
        pair_loc1[1] = get_location(path1, ts + 1)
        pair_loc2[0] = get_location(path2, ts)
        pair_loc2[1] = get_location(path2, ts + 1)
        # 检测顶点碰撞
        if pair_loc1[0] == pair_loc2[0]:
            vertex_collision = { 'ts': ts, 'loc': [pair_loc1[0]] }
            return vertex_collision
        # 检测边碰撞
        if (pair_loc1[0] == pair_loc2[0] and pair_loc1[1] == pair_loc2[1]) or (pair_loc1[0] == pair_loc2[1] and pair_loc1[1] == pair_loc2[0]):
            edge_collision = { 'ts': ts + 1, 'loc': pair_loc1 }
            return edge_collision
    return None


def detect_collisions(paths):
    # 检测所有机器人路径对之间的第一个碰撞
    agents_list = []
    for i in range(len(paths)):
        agents_list.append(i)
    # 生成所有路径对
    pairs_list = [(a, b) for idx, a in enumerate(paths) for b in paths[idx + 1:]]
    agents = [(a, b) for idx, a in enumerate(agents_list) for b in agents_list[idx + 1:]]
    first_collisions = []
    for i in range(len(pairs_list)):
        a = detect_collision(pairs_list[i][0], pairs_list[i][1])
        if a:
            collision = { 'a1': agents[i][0], 'a2': agents[i][1], 'loc': a['loc'], 'timestep': a['ts'] }
            first_collisions.append(collision)
    return first_collisions


def standard_splitting(collision):
    # 生成约束以解决给定的碰撞
    # 顶点碰撞：第一个约束阻止第一个代理在指定时间步到达指定位置，第二个约束阻止第二个代理在指定时间步到达指定位置
    # 边碰撞：第一个约束阻止第一个代理在指定时间步通过指定边，第二个约束阻止第二个代理在指定时间步通过指定边
    if len(collision['loc']) > 1:
        constraint_1 = { 'agent': collision['a1'], 'loc': collision['loc'], 'timestep': collision['timestep'], 'positive': False }
        constraint_2 = { 'agent': collision['a2'], 'loc': collision['loc'], 'timestep': collision['timestep'], 'positive': False }
        new_constraints = [constraint_1, constraint_2]
    else:
        constraint_1 = { 'agent': collision['a1'], 'loc': collision['loc'], 'timestep': collision['timestep'], 'positive': False }
        constraint_2 = { 'agent': collision['a2'], 'loc': collision['loc'], 'timestep': collision['timestep'], 'positive': False }
        new_constraints = [constraint_1, constraint_2]
    return new_constraints


def disjoint_splitting(collision):
    # 生成约束以解决给定的碰撞
    # 顶点碰撞：随机选择一个代理，第一个约束强制该代理在指定时间步到达指定位置，第二个约束阻止该代理在指定时间步到达相同位置
    # 边碰撞：随机选择一个代理，第一个约束强制该代理在指定时间步通过指定边，第二个约束阻止该代理在指定时间步通过相同边
    rand_int = random.randint(0, 1)
    rand_agent = collision['a1'] if rand_int == 0 else collision['a2']
    if len(collision['loc']) > 1:
        constraint_1 = { 'agent': rand_agent, 'loc': collision['loc'], 'timestep': collision['timestep'], 'positive': True }
        constraint_2 = { 'agent': rand_agent, 'loc': collision['loc'], 'timestep': collision['timestep'], 'positive': False}
        new_constraints = [constraint_1, constraint_2]
    else:
        constraint_1 = { 'agent': rand_agent, 'loc': collision['loc'], 'timestep': collision['timestep'], 'positive': True }
        constraint_2 = { 'agent': rand_agent, 'loc': collision['loc'], 'timestep': collision['timestep'], 'positive': False }
        new_constraints = [constraint_1, constraint_2]
    return new_constraints


def paths_violate_constraint(constraint, paths):
    # 返回违反给定约束的所有代理的 ID
    rst = []
    for i in range(len(paths)):
        if i == constraint['agent']:
            continue
        curr = get_location(paths[i], constraint['timestep'])
        prev = get_location(paths[i], constraint['timestep'] - 1)
        if len(constraint['loc']) == 1:  # 顶点约束
            if constraint['loc'][0] == curr:
                rst.append(i)
        else:  # 边约束
            if constraint['loc'][0] == prev or constraint['loc'][1] == curr or constraint['loc'] == [curr, prev]:
                rst.append(i)
    return rst


class CBSSolver(object):
    """CBS 算法的高层面搜索类"""

    def __init__(self, my_map, starts, goals):
        # 初始化 CBS 求解器
        # my_map   - 表示障碍物位置的二维列表
        # starts      - 每个代理的起始位置列表
        # goals       - 每个代理的目标位置列表
        self.my_map = my_map
        self.starts = starts
        self.goals = goals
        self.num_of_agents = len(goals)
        self.num_of_generated = 0
        self.num_of_expanded = 0
        self.CPU_time = 0
        self.open_list = []
        # 计算每个代理的启发式函数
        self.heuristics = []
        for goal in self.goals:
            self.heuristics.append(compute_heuristics(my_map, goal))

    def push_node(self, node):
        # 将节点推入优先队列
        heapq.heappush(self.open_list, (node['cost'], len(node['collisions']), self.num_of_generated, node))
        self.num_of_generated += 1

    def pop_node(self):
        # 从优先队列中弹出优先级最高的节点
        _, _, id, node = heapq.heappop(self.open_list)
        self.num_of_expanded += 1
        return node


    def find_solution(self, disjoint=True):
        # 寻找所有代理从起始位置到目标位置的路径
        self.start_time = timer.time()
        # 生成根节点
        root = {'cost': 0,
                'constraints': [],
                'paths': [],
                'collisions': []}
        # 为每个代理找到初始路径
        for i in range(self.num_of_agents):
            path = a_star(self.my_map, self.starts[i], self.goals[i], self.heuristics[i], i, root['constraints'])
            if path is None:
                raise BaseException('No solutions')
            root['paths'].append(path)
        root['cost'] = get_sum_of_cost(root['paths'])
        root['collisions'] = detect_collisions(root['paths'])
        self.push_node(root)
        # 高层面搜索循环
        while len(self.open_list) > 0:
            P = self.pop_node()
            # 如果当前节点没有碰撞，返回解决方案
            if not P['collisions']:
                self.print_results(P)
                return P['paths']
            # 选择第一个碰撞并生成约束
            collision = P['collisions'][0]
            if disjoint:
                constraints = disjoint_splitting(collision)
            else:
                constraints = standard_splitting(collision)
            # 为每个约束生成子节点
            for constraint in constraints:
                update = True
                # 处理负约束
                if constraint['positive'] == False:
                    new_constraint_set = P['constraints'].copy()
                    new_constraint_set.append(constraint)
                    Q = {'cost': 0,
                        'constraints': new_constraint_set,
                        'paths': P['paths'],
                        'collisions': []}
                    a = constraint['agent']
                    path = a_star(self.my_map, self.starts[a], self.goals[a], self.heuristics[a], a, Q['constraints'])
                    if path:
                        new_path_set = P['paths'].copy()
                        new_path_set[a] = path
                    else:
                        update = True
                # 处理正约束
                else:
                    new_constraint_set = P['constraints'].copy()
                    new_constraint_set.append(constraint)
                    new_path_set = P['paths'].copy()
                    Q = {'cost': 0,
                        'constraints': new_constraint_set,
                        'paths': P['paths'],
                        'collisions': P['collisions']}
                    a = constraint['agent']
                    path = a_star(self.my_map, self.starts[a], self.goals[a], self.heuristics[a], a, new_constraint_set)
                    if path:
                        new_path_set[a] = path
                    # 找到所有违反新约束的代理
                    agent_ids = paths_violate_constraint(constraint, P['paths'])
                    for x in agent_ids:
                        new_constraint_set = P['constraints'].copy()
                        constraint['positive'] = False
                        constraint['agent'] = x
                        new_constraint_set.append(constraint)
                        Q = {'cost': 0,
                            'constraints': new_constraint_set,
                            'paths': P['paths'],
                            'collisions': P['collisions']}
                        path = a_star(self.my_map, self.starts[x], self.goals[x], self.heuristics[x], x, new_constraint_set)
                        if path:
                            new_path_set[x] = path
                        else:
                            update = False
                Q['paths'] = new_path_set
                Q['collisions'] = detect_collisions(Q['paths'])
                Q['cost'] = get_sum_of_cost(Q['paths'])
                if update == True:
                    self.push_node(Q)
        return None


    def print_results(self, node):
        # 打印解决方案的结果
        print("\n Found a solution! \n")
        CPU_time = timer.time() - self.start_time
        print("CPU time (s):    {:.2f}".format(CPU_time))
        print("Sum of costs:    {}".format(get_sum_of_cost(node['paths'])))
        print("Expanded nodes:  {}".format(self.num_of_expanded))
        print("Generated nodes: {}".format(self.num_of_generated))